package com.dragon.english_practice_back.advisor;

import com.dragon.english_practice_back.service.impl.TokenUsageServiceImpl;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.client.ChatClientMessageAggregator;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;

@Component
@Slf4j
public class TokenRecordAdvisor implements CallAdvisor, StreamAdvisor {
    private static final String SERVICE_USER_ID = "service_user_id";

    TokenUsageServiceImpl tokenUsageServiceImpl;

    public TokenRecordAdvisor(TokenUsageServiceImpl tokenUsageServiceImpl) {
        this.tokenUsageServiceImpl = tokenUsageServiceImpl;
    }

    @NotNull
    @Override
    public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
        // 获取用户id
        Map<String, Object> content = chatClientRequest.context();
        Integer serviceUserID = content.containsKey(SERVICE_USER_ID) ? Integer.parseInt(content.get(SERVICE_USER_ID).toString()): null;

        ChatClientResponse advisedResponse = callAdvisorChain.nextCall(chatClientRequest);
        this.saveTokenUsage(serviceUserID, advisedResponse);
        return advisedResponse;
    }

    @NotNull
    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain streamAdvisorChain) {
        // 获取用户id
        Map<String, Object> content = chatClientRequest.context();
        Integer serviceUserID = content.containsKey(SERVICE_USER_ID) ? Integer.parseInt(content.get(SERVICE_USER_ID).toString()): null;

        // 获取AI执行后的token消耗量
        Flux<ChatClientResponse> advisedResponseFlux = streamAdvisorChain.nextStream(chatClientRequest);
        return  new ChatClientMessageAggregator().aggregateChatClientResponse(advisedResponseFlux, new Consumer<ChatClientResponse>() {
            @Override
            public void accept(ChatClientResponse chatClientResponse) {
                if (chatClientResponse.chatResponse() != null &&
                        chatClientResponse.chatResponse().hasFinishReasons(Set.of("STOP"))) {
                    saveTokenUsage(serviceUserID, chatClientResponse);;

                }
            }
        });
    }


    @NotNull
    @Override
    public String getName() {
        return "Token record advisor";
    }

    /**
     * 定义TokenRecordAdvisor的order为int类型的最大值，表示该advisor在chain中距离执行AI服务最远
     * @return 1000
     */
    @Override
    public int getOrder() {
        return 1000;
    }

    private void saveTokenUsage(Integer serviceUserID, ChatClientResponse chatClientResponse){
        log.debug("AFTER: {}", chatClientResponse);
        Usage usage = null;
        String model = null;
        if (chatClientResponse.chatResponse() != null) {
            usage = chatClientResponse.chatResponse().getMetadata().getUsage();
            model = chatClientResponse.chatResponse().getMetadata().getModel();
        }
        // 记录token消耗量
        if (usage != null) {
            tokenUsageServiceImpl.chatServiceUsage(
                    usage.getTotalTokens(),
                    usage.getPromptTokens(),
                    usage.getCompletionTokens(),
                    serviceUserID,
                    model
            );
        }
    }


}
